[Rust] Amazon Bedrock のナレッジベースでPinecone 使ってRAGる

Introduction

Knowledge base for Amazon Bedrock は、Amazon Bedrock で使える、
LLM と組み合わせて手軽に RAG を構築することが可能な Bedrock の機能です。
* RAGについてはこのあたりをご確認ください

今回はAmazon BedrockとベクトルデータベースのPinecondeをつかって
RAGのデモを試してみます。

Environments

• MacBook Pro (13-inch, M1, 2020)
• OS : MacOS 14.3.1
• Rust : 1.77.2

AWS アカウントはセットアップ済みとします。

Setup

では、このあたりを参考にセットアップしていきます。
今回は Knowledge base for Amazon Bedrock 機能を使って
S3にアップロードしたPDFをベクトルデータベースのPineconeに取り込みます。
そして AWS SDK for Rust をつかってプログラムから
アクセスしてみます。

具体的にはPDF形式のクラスメソッド就業規則(実際はGenAIで作成したサンプル)を
S3にアップロードし、それをPineconeに取り込んでBedrockに
「クラスメソッド株式会社の休職について教えて」と尋ねると、
AIがクラスメソッドにおける休職の規則について教えてくれるような
RAGのデモを作成しました。

Pinecone のセットアップ

PineconeをBedrockのナレッジベースにする手順はググればいくつもでてくるので
詳細は端折ります。
アカウントがない人はPineconeのサイト右上の
Sign UpSign Up Free からアカウント登録しましょう。
Amazon Marketplace からPineconeをSubbscribeしてもOKです。
私はPineconeのサイトから Googleアカウントでサインアップしました。

ダッシュボードから Create Index をクリックし、
適当なNameとDimensionには1536を入力します。
※埋め込みのTitan G1 Embeddings–Text用

あとは Capacity Mode を「Serverless」、リージョンを「us-east-1」にすれば OK です。
(今回使用するリージョンはすべて us-east-1 です)
Indexが作成されたら、下記情報をメモしておきましょう。

• HOST
• API Key

AWS Secrets Manager で API Key を登録

AWSコンソールでSecrets Managerへアクセスし、
シークレットのタイプは「その他のシークレットのタイプ」、
キーは「apiKey」値はさきほどのAPI Keyの値を設定します。
あとは適当にシークレットの名前を設定して保存すれば OK です。
詳細画面で「シークレットのARN」がみれるので、これをメモしておきます。

S3バケットの作成とファイルアップロード

AWSコンソールでS3にアクセスし、バケットを新しく作成します。
作成したらここでサポートされている形式の取り込みたいファイルをアップロードしましょう。
ここではPDF形式のクラスメソッド就業規則(Gen AI で作成したサンプル)
をアップロードしました。 ここではバケットのURIをメモしておきます。

Bedrock でナレッジベースの作成

AWSコンソールでBedrockを選択し、
左のメニューからナレッジベースをクリックします。
「ナレッジベースを作成」ボタンをクリックしてナレッジベース名を適当に設定しましょう。
データソースは先程のS3 URIを設定します。
その後埋め込みモデルを選択します。
ここでは「Titan Embeddings G1–Text」を選択します。
使用する埋め込みモデルは、Amazon Bedrock でアクセス権が付与されている必要があります。
アクセス権がないと、この後データ同期するときにエラーになります。
もしアクセス権がないなら、Bedrockのモデルアクセス画面を開いて
アクセス権のリクエストをしましょう。
「アクセスが付与されました」となっていればOKです。

あとは下記情報を設定してナレッジベース作成ボタンをクリックします。

• エンドポイントURLにPineconeのHOSTを設定
• 認証情報シークレットARNに、Secrets Manager でメモしたARNを設定
• テキストフィールドに「text」を設定
• Bedrock マネージドメタデータフィールドに「metadata」を設定

ナレッジベースのデータソース欄で Sync(同期）ボタンをクリックすると
データ同期が実行されます。
このとき ↓ のようなエラーがでた場合、
モデルのアクセス権がない可能性があるので確認してください。

ValidationException
Knowledge base role arn:aws:iam::xxxxxxxx:role/service-role/AmazonBedrockExecutionRoleForKnowledgeBase is
not able to call specified embedding model arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1:
You don't have access to the model with the specified model ID. (Service: BedrockRuntime, Status Code: 403, Request ID: xxxxxxx)

※同期ボタンクリックで何も反応がない場合、画面上部で↑のエラーが発生している可能性があります

同期が完了したらRustのプログラムを作成する準備は完了です。

Try

では Cargoでプロジェクトを作成してBedrockにアクセスしてみましょう。
この記事を参考にRust用プロジェクトを作成します。

↑ の crate と合わせて下記 crate も追加しましょう。

% cargo add aws-sdk-bedrockagent
% cargo add aws-sdk-bedrockagentruntime

ではまず、Pineconeからから取得された情報をそのまま表示してみます。
KnowledgeBaseQueryに問い合わせたい内容、
set_knowledge_base_id では、「ナレッジベースの概要」で
確認できる「ナレッジベース ID」を設定します。
また、KnowledgeBaseVectorSearchConfigurationで
set_number_of_results を設定しています。
こうすると検索結果が指定した件数になります。

async fn retrieve_knowledge_base(prompt: &str, config: &aws_config::SdkConfig) {
    let runtime = aws_sdk_bedrockagentruntime::Client::new(&config);
    let retrieve: RetrieveFluentBuilder = runtime.retrieve();
    let k_query = KnowledgeBaseQuery::builder().text(prompt).build().unwrap();

    //1件だけ取得
    let search_conf =
        KnowledgeBaseVectorSearchConfiguration::builder().set_number_of_results(Some(1));
    let kbr_conf = KnowledgeBaseRetrievalConfiguration::builder()
        .vector_search_configuration(search_conf.build());

    let retrieve_output = retrieve
        .set_knowledge_base_id(Some("<ナレッジベース ID>".to_string()))
        .set_retrieval_configuration(Some(kbr_conf.build()))
        .retrieval_query(k_query)
        .send()
        .await;

    match retrieve_output {
        Ok(output) => {
            let results = output.retrieval_results();
            for item in results {
                println!(
                    "score: {:?}, item:{:?}",
                    item.score().unwrap(),
                    item.content().unwrap().text()
                );
            }
        }
        Err(err) => match err {
            SdkError::ServiceError(err) => {
                eprintln!("Bedrock Error: {:?}", err);
            }
            SdkError::ConstructionFailure(err) => {
                eprintln!("Construction Failure: {:?}", err);
            }
            _ => {
                eprintln!("Transport Error: {:?}", err);
            }
        },
    }
}

下記のように main を記述して実行してみます。

#[tokio::main]
async fn main() -> Result<(), Error> {
    let region = aws_config::Region::new("<リージョン指定>");//us-east-1
    let config = aws_config::from_env().region(region).load().await;
    let prompt = "クラスメソッド株式会社の休職制度について教えて下さい。";
    retrieve_knowledge_base(prompt, &config).await;
    Ok(())
}

結果は下記のように、ドキュメントから関係ありそうな部分が検索されました。

% cargo run

score: 0.6958675980567932,
item:"短時間正社員への転換） 第11条 正社員は、子供の養育、家族の介護、健康上の理由などにより、短時間正社員への転換を申し出ることができる。
・・・

（休職） 第12条 従業員が、業務外の傷病、精神的疾患、公職就任・立候補、会社命令による会社外職務従事、刑事事件での起訴、その他やむを得ない事情により休職となる場合がある。
・・・

（服務の基本心得） 第13条 会社は社会的存在であり、従業員は社会人として社会的ルールとマナーを守らなければならない。
・・・

次に応答生成を有効にして実行してみます。
先ほどと違うのは Retrieve〜だったのが、RetrieveAndGenerate〜になってるところです。
応答生成には Claude3 Sonnet を使用します。
set_model_arn にモデルの ARN を設定することで、
プロンプトに対する応答をうまいことやってくれます。

async fn retrieve_knowledge_base_gen(prompt: &str, config: &aws_config::SdkConfig) {
    let runtime = aws_sdk_bedrockagentruntime::Client::new(&config);
    let retrieve: RetrieveAndGenerateFluentBuilder = runtime.retrieve_and_generate();

    let rg_conf = KnowledgeBaseRetrieveAndGenerateConfigurationBuilder::default()
        .set_knowledge_base_id(Some("<ナレッジベース ID>".to_string()))
        //モ応答生成に使用するモデルIDを設定
        .set_model_arn(Some(
            "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0"
                .to_string(),
        ))
        .build()
        .unwrap();

    let rgc = RetrieveAndGenerateConfigurationBuilder::default()
        .set_type(Some(RetrieveAndGenerateType::KnowledgeBase))
        .set_knowledge_base_configuration(Some(rg_conf))
        .build()
        .unwrap();

    let rag_input = RetrieveAndGenerateInputBuilder::default()
        .set_text(Some(prompt.to_string()))
        .build()
        .unwrap();

    let retrieve_output = retrieve
        .retrieve_and_generate_configuration(rgc)
        .set_input(Some(rag_input))
        .send()
        .await;

    match retrieve_output {
        Ok(output) => {
            println!("Result: {:?}", output.output().unwrap().text());
        }
        Err(err) => match err {
            SdkError::ServiceError(err) => {
                eprintln!("Bedrock Error: {:?}", err);
            }
            SdkError::ConstructionFailure(err) => {
                eprintln!("Construction Failure: {:?}", err);
            }
            _ => {
                eprintln!("Transport Error: {:?}", err);
            }
        },
    }
}

実行してみると ↓ のような応答がかえってきます。
しっかりとナレッジベースの情報を参照し、自然な形式で回答されました。

% cargo run

Result: "クラスメソッド株式会社では、従業員が以下のような事情により働くことが難しい場合、一定期間の休職を認めています。\n\n- 業務以外での病気やケガで2週間以上欠勤し、まだ回復していない場合
\n\n- メンタルヘルスの問題で十分な働きができない場合
\n\n- 公職に就くため会社の業務と両立できない場合
\n\n- 会社の指示で外部の仕事に従事する場合
\n\n- 刑事事件で起訴された場合
\n\n- その他、本人の事情で会社が休職が適切だと判断した場合
休職期間は、勤続年数に応じて設定されており、最長で勤続3年未満は6ヶ月、勤続3年以上は1年となります。休職中は基本的に給与は支払われず、勤続年数にもカウントされませんが、会社の命令で外部出向する場合は例外的な取り扱いがあります。
休職者は毎月1回以上、所属長に状況を報告しなければなりません。休職の理由がなくなった場合は、元の仕事に復帰しますが、元の仕事に戻ることが難しい場合は、別の仕事に就くこともあります。"

Summary

今回は Bedrock のナレッジベース機能を Pinecone でやってみました。
データソースの取り込みが S3バケットにおくだけでできますし、
単にファイルを取り込んだだけでもそれなりの応答がかえってきました。
Bedtrockの機能を使えば RAG システム構築が捗ります。

References

• LangChain + Amazon Bedrock + Momento で構築する、AWS フルサーバーレス RAG 構成のすすめ
• RAG の解説: LLM とベクトルデータベースを活用したアプローチのまとめ
• Amazon Bedrock の Knowledge base を試す
• [やってみた]Amazon Bedrock の Knowledge base を触ってみた #AWSreInvent
• Qdrant と Rust を使用した RAG Web サービスの構築
• Amazon Bedrock でモデルをカスタムして偉大なミュージシャンを降臨させた(?)話
• 生成 AI のトレンド「LLM+RAG」を解説
• [アップデート] Amazon Bedrock で新モデル「Cohere Command R/R+」が利用可能になったので、RAG で使ってみた